#!/usr/bin/env python3
"""
模型导出脚本 - 将PyTorch模型文件(.pt)转换为TensorRT引擎格式(.engine)

使用方法:
    python export_model.py /path/to/model.pt
    python export_model.py /path/to/model.pt --imgsz 640 --device 0
"""

import argparse
import os
import sys
from pathlib import Path

try:
    from ultralytics import YOLO
except ImportError:
    print("错误: 请先安装ultralytics库")
    print("运行: pip install ultralytics")
    sys.exit(1)


def export_to_engine(pt_path, imgsz=640, device=0, half=False, batch=1, nms=False, data=None, 
                     dynamic=False, int8=False, fraction=1.0, workspace=4, 
                     conf=0.25, iou=0.45, max_det=100, verbose=True):
    """
    将PT模型文件导出为TensorRT引擎格式
    
    Args:
        pt_path (str): PT模型文件路径
        imgsz (int): 输入图像尺寸
        device (int): GPU设备ID
        half (bool): 是否使用FP16精度
        batch (int): 批处理大小
        nms (bool): 是否将NMS融合到导出的模型中
        data (str): 数据集配置文件路径，用于INT8校准
        dynamic (bool): 是否启用动态形状
        int8 (bool): 是否使用INT8量化
        fraction (float): 用于INT8校准的数据集比例(0.0-1.0)
        workspace (float): TensorRT工作空间大小(GiB)
        conf (float): NMS置信度阈值
        iou (float): NMS IOU阈值
        max_det (int): 最大检测数量
        verbose (bool): 是否显示详细信息
    
    Returns:
        str: 导出的引擎文件路径
    """
    # 检查PT文件是否存在
    pt_path = Path(pt_path)
    if not pt_path.exists():
        raise FileNotFoundError(f"PT文件不存在: {pt_path}")
    
    if not pt_path.suffix.lower() == '.pt':
        raise ValueError(f"文件必须是.pt格式: {pt_path}")
    
    # 获取输出文件路径
    engine_path = pt_path.with_suffix('.engine')
    
    if verbose:
        print(f"输入文件: {pt_path}")
        print(f"输出文件: {engine_path}")
        print(f"图像尺寸: {imgsz}")
        print(f"批处理大小: {batch}")
        print(f"设备: cuda:{device}")
        print(f"半精度: {half}")
        print(f"INT8量化: {int8}")
        print(f"动态形状: {dynamic}")
        print(f"融合NMS: {nms}")
        if nms:
            print(f"  - 置信度阈值: {conf}")
            print(f"  - IOU阈值: {iou}")
            print(f"  - 最大检测数: {max_det}")
        print(f"工作空间: {workspace}GiB")
        if data:
            print(f"数据集配置: {data}")
        if int8:
            print(f"校准数据比例: {fraction}")
        print("-" * 50)
    
    try:
        # 加载YOLO模型
        if verbose:
            print("正在加载模型...")
        model = YOLO(str(pt_path))
        
        # 导出为TensorRT引擎
        if verbose:
            print("开始导出TensorRT引擎...")
        
        exported_path = model.export(
            format='engine',
            imgsz=imgsz,
            device=device,
            half=half,
            batch=batch,
            nms=nms,
            data=data,
            dynamic=dynamic,
            int8=int8,
            fraction=fraction,
            workspace=workspace,
            conf=conf,
            iou=iou,
            max_det=max_det,
            verbose=verbose
        )
        
        # 移动文件到PT文件所在目录（如果需要）
        exported_path = Path(exported_path)
        if exported_path.parent != pt_path.parent:
            target_path = pt_path.parent / exported_path.name
            if target_path.exists():
                if verbose:
                    print(f"覆盖已存在的文件: {target_path}")
                target_path.unlink()
            
            exported_path.rename(target_path)
            exported_path = target_path
        
        if verbose:
            print(f"✅ 导出成功: {exported_path}")
            print(f"文件大小: {exported_path.stat().st_size / 1024 / 1024:.2f} MB")
        
        return str(exported_path)
        
    except Exception as e:
        print(f"❌ 导出失败: {str(e)}")
        raise


def main():
    parser = argparse.ArgumentParser(
        description="将PyTorch模型文件(.pt)转换为TensorRT引擎格式(.engine)",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
示例:
  python export_model.py model.pt
  python export_model.py /path/to/model.pt --imgsz 640 --device 0
  python export_model.py model.pt --imgsz 1280 --device 0 --no-half
  python export_model.py model.pt --batch 8 --dynamic --workspace 8
  python export_model.py model.pt --int8 --data dataset.yaml --fraction 0.5
  python export_model.py model.pt --nms --dynamic --workspace 6
        """
    )
    
    parser.add_argument(
        '--pt_path',
        type=str,
        help='PT模型文件路径'
    )
    
    parser.add_argument(
        '--imgsz',
        type=int,
        default=640,
        help='输入图像尺寸 (默认: 640)'
    )
    
    parser.add_argument(
        '--device',
        type=int,
        default=0,
        help='GPU设备ID (默认: 0)'
    )
    
    parser.add_argument(
        '--no-half',
        action='store_true',
        help='禁用FP16半精度模式'
    )
    
    parser.add_argument(
        '--batch',
        type=int,
        default=1,
        help='批处理大小 (默认: 1)'
    )
    
    parser.add_argument(
        '--nms',
        action='store_true',
        help='将NMS融合到导出的模型中'
    )
    
    parser.add_argument(
        '--data',
        type=str,
        default=None,
        help='数据集配置文件路径，用于INT8校准'
    )
    
    parser.add_argument(
        '--dynamic',
        action='store_true',
        help='启用动态形状支持'
    )
    
    parser.add_argument(
        '--int8',
        action='store_true',
        help='使用INT8量化'
    )
    
    parser.add_argument(
        '--fraction',
        type=float,
        default=1.0,
        help='用于INT8校准的数据集比例 (0.0-1.0，默认: 1.0)'
    )
    
    parser.add_argument(
        '--workspace',
        type=float,
        default=8,
        help='TensorRT工作空间大小(GiB) (默认: 8)'
    )
    
    parser.add_argument(
        '--quiet',
        action='store_true',
        help='静默模式，减少输出信息'
    )
    
    parser.add_argument(
        '--conf', 
        type=float, 
        default=0.25,
        help='NMS 置信度阈值（仅 nms 模式，默认: 0.25）'
    )
    parser.add_argument(
        '--iou', 
        type=float, 
        default=0.45,
        help='NMS IOU 阈值（仅 nms 模式，默认: 0.45）'
    )
    parser.add_argument(
        '--max-det', 
        type=int, 
        default=100,
        help='最大检测数量（仅 nms 模式，默认: 100）'
    )
    
    args = parser.parse_args()
    
    try:
        # 导出模型
        engine_path = export_to_engine(
            pt_path=args.pt_path,
            imgsz=args.imgsz,
            device=args.device,
            half=not args.no_half,
            batch=args.batch,
            nms=args.nms,
            data=args.data,
            dynamic=args.dynamic,
            int8=args.int8,
            fraction=args.fraction,
            workspace=args.workspace,
            conf=args.conf,
            iou=args.iou,
            max_det=args.max_det,
            verbose=not args.quiet
        )
        
        if not args.quiet:
            print(f"\n🎉 模型导出完成!")
            print(f"引擎文件: {engine_path}")
        
    except KeyboardInterrupt:
        print("\n⚠️  用户中断操作")
        sys.exit(1)
    except Exception as e:
        print(f"\n❌ 错误: {str(e)}")
        sys.exit(1)


if __name__ == "__main__":
    main()
